import json
import xml.etree.ElementTree as ET

"""
参考: https://blog.csdn.net/qq_15969343/article/details/80848175
https://blog.csdn.net/qq_41375609/article/details/94737915?depth_1-utm_source=distribute.pc_relevant.none-task&utm_source=distribute.pc_relevant.none-task
coco数据集介绍
{
    "info"        : info,
    "images"      : [],
    "annotations" : [annotations],
    "license"     : [license],
    "categories"  : [category]
}
其中info是一个字典
info = {
    "year":int,
    "version" : int,
    "description" : str,
    'contributor': str,
    'url': str,
    'date_created': datetime
}
images 是一个列表，每一个元素是一个字典，存放image的信息
annotations是一个列表，每一个元素是一个图片的一个box的标注的信息，
"""

label_to_id = {'zhujia': 1, 'fujia': 2}


def get_default_dict():
    default_dict = {'categories': [{
        "supercategory": "driver",
        "id": 1,
        "name": "zhujia"
    }, {
        "supercategory": "driver",
        "id": 2,
        "name": "fujia"
    }],
        'images': [],
        'annotations': [],
        'info': {
            'year': 2020,
            'version': 'v1',
            'description': 'driver_detect',
            'contributor': 'houkai',
            'url': None,
            'date_created': '2020-03-25'
        },
        'license': {
            "url": "no url",
            "id": 1,
            "name": "Attribution-NonCommercial-ShareAlike License"
        }}

    return default_dict


def parse_single_xml(image_id, image_path, annotation_path):
    tree = ET.parse(annotation_path)
    root = tree.getroot()
    objs = root.findall('object')
    image_width = int(root.find('size').find('width').text)
    image_height = int(root.find('size').find('height').text)
    file_name = image_path
    annotations = []
    image = {'file_name': file_name,
             'id': image_id,  # 图片的ID编号（每张图片ID是唯一的）
             'width': image_width,
             'height': image_height}
    for idx, obj in enumerate(objs):
        label = obj.find('name').text
        x1 = int(obj.find('bndbox').find('xmin').text)
        y1 = int(obj.find('bndbox').find('ymin').text)
        x2 = int(obj.find('bndbox').find('xmax').text)
        y2 = int(obj.find('bndbox').find('ymax').text)
        box_width = max(0, x2 - x1)
        box_height = max(0, y2 - y1)
        annotations.append({
            'area': box_width * box_height,
            'bbox': [x1, y1, box_width, box_height],  # 定位边框 [x,y,w,h]
            'category_id': label_to_id[label],  # 对应的图片ID（与images中的ID对应）
            'id': idx + 1,  # 对象ID，因为每一个图像有不止一个标注对象，所以要对每一个对象编号（每个对象的ID是唯一的）
            'image_id': image_id,  # 对应的图片ID（与images中的ID对应）
            'iscrowd': 0,
            # mask, 矩形是从左上角点按顺时针的四个顶点
            'segmentation': [x1, y1, x2, y1, x2, y2, x1, y2],
        })

    return image, annotations


def convert_data(type):
    with open(type + '.txt', 'r') as file:
        lines = file.readlines()
        for index, line in enumerate(lines):
            image_path, xml_path = line.strip().split(' ')
            image, annotations = parse_single_xml(index, image_path, xml_path)
            save_dict['images'].append(image)
            save_dict['annotations'].extend(annotations)
            if index % 1000 == 0 and index / 1000 != 0:
                print("already process {} images".format(index))

    with open(type + '.json', 'w') as f:
        f.write(json.dumps(save_dict, ensure_ascii=False))  # ensure_ascii=False解决中文乱码


if __name__ == '__main__':
    save_dict = get_default_dict()
    # convert_data('train')
    convert_data('test')
